import argparse
from audioop import add
from tkinter.messagebox import NO

def get_basic_parser_for_train() -> argparse.ArgumentParser:
        """
        return basic arg parser
        """
        parser = argparse.ArgumentParser(description="Training")
        parser.add_argument('--model_name', type=str, required=True)
        parser.add_argument('--project', type=str, help='wandb的项目名称',required=True)
        parser.add_argument('--train_path', type=str,required=True)
        parser.add_argument('--dev_path', type=str)
        parser.add_argument('--test_path', type=str, required=True)
        parser.add_argument('--epochs', type=int, required=True)
        parser.add_argument('--steps', type=int, help='Training step for distributed training')
        parser.add_argument('--batch_size', type=int, required=True)
        parser.add_argument('--max_length', type=int, required=True)
        parser.add_argument('--weight_decay', type=float, default=0.01, required=True)
        parser.add_argument('--scheduler_name', type=str, default='linear', required=True)
        parser.add_argument('--warm_factor', type=float, required=True, default=0.1)
        parser.add_argument('--seed', type=int, required=True)
        parser.add_argument('--lm_learning_rate', type=float, required=True)
        parser.add_argument('--dropout', type=float, default=0.3)
        parser.add_argument('--output_path', type=str, required=True)
        parser.add_argument('--device_ids', type=str, help="GUP ids")
        parser.add_argument('--others_learning_rate', type=float)
        parser.add_argument('--checkpoint_save_path', type=str, default=None)
        parser.add_argument('--pretrained_checkpoint_path', type=str, default=None)
        parser.add_argument('--log_path', type=str)
        parser.add_argument('--workers', type=int, default=1)
        parser.add_argument('--classifier_dropout', type=float)
        parser.add_argument('--inference', default=False, action='store_true')
        parser.add_argument('--multiply_lr', default=False, action='store_true')
        parser.add_argument('--attack', default=False, action='store_true')
        parser.add_argument('--max_grad_norm', type=int, default=2)
        parser.add_argument('--k_shot', type=str, help='负数表示全部的数据用于训练')
        parser.add_argument('--data_name', type=str)
        parser.add_argument('--method', type=str)
        parser.add_argument('--wandb_mode', type=str, default='online', help='"online", "offline" or "disabled". Defaults to online.')
        return parser

def get_basic_parser_for_other_train() -> argparse.ArgumentParser:
        parser = get_basic_parser_for_train()
        parser.add_argument('--weight_label', nargs='+', type=float, help='BCELOSS的标签平衡的使用')
        parser.add_argument('--max_entity_len', type=int, required=True)
        parser.add_argument('--min_entity_len', type=int, required=True)
        parser.add_argument('--user_dict', type=str, required=True)
        parser.add_argument('--pet_patterns_idx', type=int, default=-1)
        parser.add_argument('--calculation_by_category', default=True, action='store_false')
        parser.add_argument('--nest', default=False, action='store_true')
        parser.add_argument('--futher_segmentation', default=False, action='store_true')
        parser.add_argument('--demonstration', default=False, action='store_true')
        parser.add_argument('--predict_step', type=int, default=50, help="在推理阶段用的step数量，如果用全部太慢了")
        return parser

def get_basic_parser_for_inference() -> argparse.ArgumentParser:
        parser = argparse.ArgumentParser(description="Inference")
        parser.add_argument('--save_path', type=str, default='./model.pt')
        parser.add_argument('--dataset_path', type=str)
        parser.add_argument('--test_dir', type=str)
        parser.add_argument('--predict_path', type=str, default='./output.json')
        parser.add_argument('--output_path', type=str, default=None)
        parser.add_argument('--dev_path', type=str, default=None)
        parser.add_argument('--log_path', type=str, default=None)
        parser.add_argument('--dist_emb_size', type=int)
        parser.add_argument('--type_emb_size', type=int)
        parser.add_argument('--lstm_hid_size', type=int)
        parser.add_argument('--conv_hid_size', type=int)
        parser.add_argument('--workers', type=int,default=1)
        parser.add_argument('--bert_hid_size', type=int)
        parser.add_argument('--ffnn_hid_size', type=int)
        parser.add_argument('--biaffine_size', type=int)
        parser.add_argument('--dilation', nargs='+', type=int)

        parser.add_argument('--emb_dropout', type=float)
        parser.add_argument('--conv_dropout', type=float)
        parser.add_argument('--out_dropout', type=float)

        parser.add_argument('--epochs', type=int)
        parser.add_argument('--batch_size', type=int)

        parser.add_argument('--clip_grad_norm', type=float)
        parser.add_argument('--learning_rate', type=float)
        parser.add_argument('--weight_decay', type=float)

        parser.add_argument('--bert_name', type=str)
        parser.add_argument('--bert_learning_rate', type=float)
        parser.add_argument('--warm_factor', type=float)
        parser.add_argument('--use_bert_last_4_layers', default=False, action='store_true')
        parser.add_argument('--seed', type=int)
        return parser